# coding: utf-8

from __future__ import absolute_import, division, print_function, unicode_literals

import os
from string import Template

from apitest.cmd.generate import GenCode
from apitest.cmd.output import OutFileType, out_file_path
from apitest.cmd.thrift2pyi import Thrift2pyi


class RPCClientCodeTemplate:
    _service_string = """from typing import Union

import thriftpy2.protocol.json as proto
from apitest.engine.client.rpc import RpcClient
from apitest.engine.const.enums import RType

from api.$psm_name.$idl_name import *

PSM = "$psm"


class $client_class(RpcClient):

    def __init__(self, timeout=10, **kwargs):
        super().__init__($service_class, PSM, timeout=timeout, **kwargs)
"""
    _method_string = """
    def $method_name(self, req: Union[$req, dict], rtype: RType = RType.obj):
        if isinstance(req, dict):
            req = proto.struct_to_obj(req, $req())
        resp: $resp = self.client.$method_name(req)
        if rtype == RType.json:
            return proto.struct_to_json(resp)
        return resp
"""

    @property
    def service_string(self):
        return self._service_string

    @property
    def method_string(self):
        return self._method_string


class GenRPCCode(GenCode):

    def __init__(self, psm, idl, out_dir):
        super(GenRPCCode, self).__init__(psm=psm, idl=idl, out_dir=out_dir)
        self.thrift.parse_service()
        self.idl_name = os.path.basename(self.thrift.rm_suffix(self.idl)).replace(".thrift", "_thrift")
        self.service_mate = self.thrift.service_mate

    def generate_rpc_client(self):
        out = out_file_path(self.out_dir, OutFileType.rpc_client)
        client_tmp = RPCClientCodeTemplate()
        client_content = ""
        psm_name = self.psm.replace(".", "_")
        for service, values in self.service_mate.items():
            service_tmp_string = Template(client_tmp.service_string)
            client_class = service.replace("Service", "Client")
            client_content += service_tmp_string.substitute(
                {"psm_name": psm_name, "idl_name": self.idl_name, "psm": self.psm, "service_class": service,
                 "client_class": client_class}
            )
            for value in values:
                client_content += Template(client_tmp.method_string).substitute(
                    {"method_name": value["method_name"], "req": value["req"], "resp": value["resp"]}
                )
        with open(out, "w", encoding='utf8') as f:
            f.write(client_content)

    def generate_rpc_model(self):
        self.thrift.parse_idl_includes(self.idl, self.out_dir)
        Thrift2pyi(os.path.join(self.out_dir, os.path.basename(self.thrift.rm_suffix(self.idl))), prefix="",
                   out=""
                   ).output()
